/*****************************************************************************/
/*                                                                           */
/*    Domino Operation System System Specific Module                         */
/*                                                                           */
/*    Copyright (C) 2010 Laszlo Arvai                                        */
/*                                                                           */
/*    ------------------------------------------------------------------     */
/*    netDNS - Domain Name Service client functions                          */
/*****************************************************************************/

/*****************************************************************************/
/* Includes                                                                  */
/*****************************************************************************/
#include <krnlSystemTimer.h>
#include <netDNS.h>
#include <netIP.h>

/*****************************************************************************/
/* Cconst                                                                    */
/*****************************************************************************/
#define netDNS_RESOLVER_BUFFER_LENGTH 256
#define netDNS_PORT										 53u	// default DNS port
#define netDNS_TIMEOUT								1000	// DNS query timeout in ms

#define netDNS_TYPE_A            (1u)				// Constant for record type in DNSResolve.  Indicates an A (standard address) record.
#define netDNS_TYPE_MX           (15u)      // Constant for record type in DNSResolve.  Indicates an MX (mail exchanger) record.

/*****************************************************************************/
/* Types                                                                     */
/*****************************************************************************/

// State machine for a DNS query
typedef enum
{
   DNS_START = 0,             // Initial state to reset client state variables
   DNS_QUERY,									// Send DNS query to DNS server
   DNS_GET_RESULT,            // Wait for response from DNS server
   DNS_FAIL,									// ARP or DNS server not responding
   DNS_DONE										// DNS query is finished
} netDNSState;

/*****************************************************************************/
/* Local function prototypes                                                 */
/*****************************************************************************/
static dosWord netDNSDiscardName(dosByte* buffer, dosWord in_pos, dosWord in_length);
static void netDNSResolver(void);

/*****************************************************************************/
/* Global variables                                                          */
/*****************************************************************************/
IN_ADDR g_netPrimaryDNSAddress;
IN_ADDR g_netSecondaryDNSAddress;

/*****************************************************************************/
/* Module local variables                                                    */
/*****************************************************************************/
static IN_ADDR l_host_ip_address = {0};
static dosBool l_host_ip_address_valid = dosFalse;
static dosChar* l_host_name = DOS_NULL;
static netDNSState l_resolver_state = DNS_DONE;
static dosByte l_attempt_count = 0;
static SOCKET l_udp_socket = INVALID_SOCKET;
static dosWord l_transaction_id = 0;
static dosByte l_query_type = netDNS_TYPE_A;
static dosWord l_time_stamp = 0;
static dosByte l_is_opened = dosFalse;

/*****************************************************************************/
/* DNS functions                                                             */
/*****************************************************************************/

///////////////////////////////////////////////////////////////////////////////
//! Opens DNS client and starts address resolution
//! \param String containing domain name
//! \return dosTrue if resolution procedure is started, dosFalse when resolver is aready opened
dosBool netDNSOpen(dosString in_name)
{
	// check if resolver is opened
	if( l_is_opened )
	{
		// refuse request if resolver is opened
		return dosFalse;
	}

	// init resolver
	l_is_opened = dosTrue;
	l_host_name = in_name;
	l_query_type = netDNS_TYPE_A;
	l_resolver_state = DNS_START;

	// success
	return dosTrue;
}

///////////////////////////////////////////////////////////////////////////////
//! Returns DNS resolver status (and calls worker function)
//! \param none
//! \return dosTrue if resolver is busy
dosBool netDNSIsBusy(void)
{
	// call resolver process if needed
	if( l_resolver_state != DNS_DONE )
	{
		netDNSResolver();
		return dosTrue;
	}
	else
		return dosFalse;
}

///////////////////////////////////////////////////////////////////////////////
//! Closes DNS resolver
void netDNSClose(void)
{
	// close socket
	if( l_udp_socket != INVALID_SOCKET )
	{
		closesocket( l_udp_socket);
		l_udp_socket = INVALID_SOCKET;
	}

	l_is_opened = dosFalse;
}

///////////////////////////////////////////////////////////////////////////////
//! Returns resolved DNS address
//! \param Pointer receiving the IP address
//! \return True if the IP address is valid
dosBool netDNSGetResolvedIPAddress(IN_ADDR* out_address)
{
  if( l_host_ip_address_valid )
	  *out_address = l_host_ip_address;
	
	return l_host_ip_address_valid;
}

/*****************************************************************************/
/* Local functions                                                           */
/*****************************************************************************/
static dosWord netDNSDiscardName(dosByte* buffer, dosWord in_pos, dosWord in_length)
{
 dosByte i;

 while(in_pos < in_length)
 {
  // Get first dosByte which will tell us if this is a 16-bit pointer or the 
  // length of the first of a series of labels
	i = buffer[in_pos++];
  
  // Check if this is a pointer, if so, get the reminaing 8 bits and return
  if((i & 0xC0u) == 0xC0u)
  {
   in_pos++;
   break;
  }

  // Exit once we reach a zero length label
  if(i == 0u)     
   break;

  // Discard complete label
  in_pos += i;  
 }
 
 return in_pos;
}

///////////////////////////////////////////////////////////////////////////////
//! DNS resolver worker function
void netDNSResolver(void)
{
	dosByte buffer[netDNS_RESOLVER_BUFFER_LENGTH];
	dosWord pos;
	dosWord length_pos;
	dosByte string_pos;
	dosChar ch;
	int length;
	struct sockaddr_in socket_address;
	int address_length = sizeof(socket_address);
	unsigned long iMode = 1;
	dosWord dosWord_buffer;
  dosWord questions;
  dosWord answers;
  dosWord authoritative_records;
  dosWord additional_records;
  dosByte field_index = 0;
  dosWord records_count;
  dosWord response_type;
  dosWord response_class;
  dosWord response_len;

	// check state
	switch( l_resolver_state )
	{
		// start DNS query
		case DNS_START:
			// init
			l_attempt_count = 0;
			l_host_ip_address.s_addr = 0;
			l_host_ip_address_valid = dosFalse;

			// fallthrough

		// send DNS query
		case DNS_QUERY:
			// open UDP socket
			if( (l_udp_socket = socket( AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == SOCKET_ERROR )
			{
				l_resolver_state = DNS_FAIL;
				break;
			}
			
			// bind socket
			socket_address.sin_family = AF_INET;
			socket_address.sin_port = 52905;//netDNS_PORT;
			socket_address.sin_addr.S_un.S_addr = IP_ADDR_ANY;
			if( bind( l_udp_socket, (struct sockaddr*)&socket_address, sizeof(socket_address) ) == SOCKET_ERROR )
			{
				l_resolver_state = DNS_FAIL;
				break;
			}

			// prepare query
      l_transaction_id++;

			buffer[ 0] = HIGH(l_transaction_id);	// User chosen transaction ID
      buffer[ 1] = LOW(l_transaction_id);
			buffer[ 2] = 0x01;							      // Standard query with recursion
      buffer[ 3] = 0x00;   
      buffer[ 4] = 0x00;										// 0x0001 questions
      buffer[ 5] = 0x01;
      buffer[ 6] = 0x00;										// 0x0000 answers
      buffer[ 7] = 0x00;
      buffer[ 8] = 0x00;										// 0x0000 name server resource records
      buffer[ 9] = 0x00;
      buffer[10] = 0x00;										// 0x0000 additional records
      buffer[11] = 0x00;
			pos = 12;

      // Put hostname string to resolve
			string_pos = 0;
			length_pos = pos++;

			while(1)
			{
	      do
		    {
					ch = l_host_name[string_pos++];
			    buffer[pos++] = ch;
				} while((ch != '\0') && (ch != '.') && (ch != '/') && (ch != ',') && (ch != '>'));
   
				// Put the length
				buffer[length_pos] = pos - length_pos - 2;
				length_pos = pos - 1;

				// exit at the end of the string
	      if(ch == '\0' || ch == '/' || ch == ',' || ch == '>')
         break;
			}

			// put terminator zero
			pos--;
			buffer[pos++] = '\0';

      buffer[pos++] = 0x00;      // Type: DNS_TYPE_A A (host address) or DNS_TYPE_MX for mail exchange
      buffer[pos++] = l_query_type;
      buffer[pos++] = 0x00;      // Class: IN (Internet)
      buffer[pos++] = 0x01;

			// store start time
      l_time_stamp = krnlGetSystemTimer();

			// prepare query
			socket_address.sin_family = AF_INET;
			socket_address.sin_port = netDNS_PORT;
			
			// select DNS server
			if( l_attempt_count == 0 )
				socket_address.sin_addr.S_un.S_addr = g_netPrimaryDNSAddress.s_addr;
			else
				socket_address.sin_addr.S_un.S_addr = g_netSecondaryDNSAddress.s_addr;
				
			// send query
			length = sendto(l_udp_socket, buffer, pos, 0,(struct sockaddr*)&socket_address, address_length);
			if( length == pos )
			{
				// packet send was success -> next stage
				l_resolver_state = DNS_GET_RESULT;
			}
			else
			{
				// packet send failed -> DNA falure
				l_resolver_state = DNS_FAIL;
			}
			break;

		// waits for DNS server response
		case DNS_GET_RESULT:
			// check for timeout
			if( krnlGetTimeSince(l_time_stamp) > netDNS_TIMEOUT )
			{
				l_resolver_state = DNS_FAIL;
				break;
			}

			// check for result
			length = recvfrom( l_udp_socket, buffer, sizeof(buffer), 0, (struct sockaddr*)&socket_address, &address_length);
			if( length > 0 )
			{
				// check socket parameters
				if( socket_address.sin_family	!= AF_INET || socket_address.sin_port	!= netDNS_PORT )
					return;
					
				// check sender IP
				if( ( l_attempt_count == 0 && socket_address.sin_addr.S_un.S_addr != g_netPrimaryDNSAddress.s_addr ) ||
						( l_attempt_count > 0 && socket_address.sin_addr.S_un.S_addr != g_netSecondaryDNSAddress.s_addr ))
					return;
			
				// Retrieve the DNS header and de-big-endian it
				pos = 0;
				dosWord_buffer = ((dosWord)buffer[pos] << 8) + buffer[pos+1];

				// Throw this packet away if it isn't in response to our last query
				if(dosWord_buffer != l_transaction_id)
				{
					break;
				}
				
				// skip transaction ID and Flags in the buffer
				pos += 2 + 2;
				
				// get records count
				questions = (buffer[pos+0] << 8) + buffer[pos+1];
				answers = (buffer[pos+2] << 8) + buffer[pos+3];
				authoritative_records = (buffer[pos+4] << 8) + buffer[pos+5];
				additional_records = (buffer[pos+6] << 8) + buffer[pos+7];
				pos += 8;
				
				// skip questions
				pos = netDNSDiscardName(buffer, pos, length);
				
				// skip question type and question class
				pos += 2 + 2;
				
				// process reaining three fields
				for( field_index = 0; field_index < 3 && pos < length && !l_host_ip_address_valid; field_index++ )
				{
					// select field
					switch( field_index )
					{
						// answers records
						case 0:
							records_count = answers;
							break;
							
						// authoritative records 
						case 1:
							records_count = authoritative_records;
							break;
						
						// additional records			
						case 2:
							records_count = additional_records;
							break;
					}
					
					// process records
					while( records_count > 0 && pos < length && !l_host_ip_address_valid )
					{
						// discard name
						pos = netDNSDiscardName(buffer, pos, length);
						
						// get response header
						response_type = (buffer[pos+0] << 8) + buffer[pos+1];
						response_class = (buffer[pos+2] << 8) + buffer[pos+3];
						response_len = (buffer[pos+8] << 8) + buffer[pos+9];
						pos += 10;
						
						// Make sure that this is a 4 dosByte IP address, response type A or MX, class 1
						// Check if this is Type A or MX
						if( response_type == 0x0001u && response_class== 0x0001u && response_len == 0x0004u)
						{
						 l_host_ip_address_valid = dosTrue;
						 l_host_ip_address.s_addr = buffer[pos+0] + (buffer[pos+1] << 8) + (buffer[pos+2]<<16) + (buffer[pos+3]<<24);
						}
						else
						{
							// skip response
							pos += response_len;
						}
					}
						
					// next record
					records_count--;
				}
					
				// set next state
				if( l_host_ip_address_valid )
					l_resolver_state = DNS_DONE;
				else
					l_resolver_state = DNS_FAIL;
			}
			break;

		// resolution is failed
		case DNS_FAIL:
			// check retry count
			if( l_attempt_count >= 1 )
			{
				// DNS resolution failed
				l_host_ip_address_valid = dosTrue;
				l_host_ip_address.s_addr = 0;
				l_resolver_state = DNS_DONE;
				break;
			}
			
			// retry address resolution using secondary DNS server
			l_resolver_state = DNS_QUERY;
			break;
   
		// DNS resolution finished
		case DNS_DONE:
			closesocket( l_udp_socket );
			break;
	}
}

